# -*- coding:utf8 -*-
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.model_selection import KFold
from sklearn.model_selection import StratifiedKFold

# 样本较多时, 直接划分训练集和校验集
X, y = np.arange(20).reshape(10, 2), range(10)
print("X is:\n", X)
print("y is: \n", y)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.33, random_state=42)
print("X_train is:\n", X_train)
print("X_test is:\n", X_test)
print("y_train is:\n", y_train)
print("y_test is:\n", y_test)


# 样本较少时, 练习KFold的方法
X = np.array([[11, 21], [31, 41], [1, 2], [3, 4], [5, 6], [7, 8], [7, 8], [7, 8], [7, 8]])
y = np.array([31, 31, 77, 77, 88, 88, 31, 77, 88])
print(X)
print(y)
skf = StratifiedKFold(n_splits=2, random_state=None, shuffle=False)
n_splits = skf.get_n_splits(X, y)
print("X,y交叉验证的次数:{}".format(n_splits))

for train_index, test_index in skf.split(X, y):
    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]
    print("训练集\n", "X\n", X_train, "\ny\n", y_train)
    print("测试集\n", "X\n", X_test, "\ny\n", y_test)
    print("")

print("\n KFold \n")
skf = KFold(n_splits=9)
count = 0
for train_index, test_index in skf.split(X, y):
    count = count + 1
    print("第{}次".format(count))

    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]
    print("训练集\n", "X\n", X_train, "\ny\n", y_train)
    print("测试集\n", "X\n", X_test, "\ny\n", y_test)
    print("")
